# tcp_clients.py
import asyncio
import json
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from typing import Dict, List, Optional

from tcp_server.websocket_manager import websocket_manager
from utils.log_util import logger


@dataclass
class TCPClient:
    """TCP客户端连接信息"""
    client_id: str
    ip_address: str
    port: int
    reader: asyncio.StreamReader
    writer: asyncio.StreamWriter
    connected_at: datetime
    last_heartbeat: datetime
    status: str = "connected"
    current_command: str = None

class EnhancedClientManager:
    """增强型客户端管理器 - 管理TCP连接和状态"""
    
    def __init__(self):
        self.clients: Dict[str, TCPClient] = {}  # client_id -> TCPClient
        self.client_messages: Dict[str, List[dict]] = {}  # 消息历史
        self.last_status_snapshot: Dict[str, dict] = {}  # ✅ 保存每个客户端的最后状态快照，用于对比变化
        self.client_user_map: Dict[str, Optional[int]] = {}  # ✅ client_id -> user_id 映射缓存
    
    async def add_or_update_client(
        self,
        client_id: str,
        ip_address: str,
        port: int,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
        status_data: dict = None,
        db_session=None
    ) -> str:
        """添加或更新客户端"""
        
        # 如果客户端已存在，更新连接信息
        if client_id in self.clients:
            client = self.clients[client_id]
            client.last_heartbeat = datetime.now()
            client.reader = reader
            client.writer = writer
            # ✅ 移除冗余日志：心跳更新客户端连接是高频操作
            # logger.debug(f"更新客户端连接: {client_id}")
        else:
            # 新客户端
            client = TCPClient(
                client_id=client_id,
                ip_address=ip_address,
                port=port,
                reader=reader,
                writer=writer,
                connected_at=datetime.now(),
                last_heartbeat=datetime.now()
            )
            self.clients[client_id] = client
            self.client_messages[client_id] = []
            logger.info(f"新客户端连接: {client_id} from {ip_address}:{port}")
        
        # 更新数据库中的客户端状态
        if db_session and status_data:
            await self._update_db_client_status(
                db_session, client_id, ip_address, port, status_data
            )
            # ✅ 更新client_user_map缓存（从数据库加载user_id）
            await self._refresh_client_user_map(db_session, client_id)
        
        # ✅ 优化：减少client_update推送频率
        # 只在以下情况推送client_update：
        # 1. 客户端首次连接
        # 2. 重要状态变化（离线->在线、登录状态变化等）
        # 避免频繁推送，因为status_update已经包含了状态信息
        # asyncio.create_task(self._notify_client_update(client_id))  # ✅ 暂时禁用，避免与status_update重复
        
        return client_id
    
    async def remove_client(self, client_id: str, db_session=None):
        """移除客户端"""
        if client_id in self.clients:
            client = self.clients[client_id]
            logger.info(f"客户端断开: {client_id}")
            
            # 从内存中删除
            del self.clients[client_id]
            
            # ✅ 清理状态快照
            if client_id in self.last_status_snapshot:
                del self.last_status_snapshot[client_id]
            
            # ✅ 清理user_id映射（但保留，因为可能后续还会连接）
            # 注意：不删除client_user_map，因为Agent可能重新连接
            
            # 更新数据库状态为离线
            if db_session:
                await self._mark_client_offline(db_session, client_id)
            
            # 通知WebSocket
            asyncio.create_task(websocket_manager.broadcast({
                "type": "client_disconnected",
                "client_id": client_id,
                "timestamp": datetime.now().isoformat()
            }))

    async def send_to_client(self, client_id: str, message: dict) -> bool:
        """发送消息到客户端"""
        client = self.get_client_by_id(client_id)
        if not client:
            logger.warning(f"⚠️ 无法发送消息，客户端 {client_id} 不存在")
            return False
        
        if client.writer.is_closing():
            logger.warning(f"⚠️ 无法发送消息，客户端 {client_id} 的writer正在关闭")
            return False
        
        try:
            # 设置当前命令状态
            event = message.get("event")
            if event:
                client.current_command = event
            
            # 调试：记录发送的详细消息（包括搜索命令）
            if event in ["cancel_order", "search"]:
                logger.info(f"📤 准备发送 {event} 命令到 {client_id}: message={message}")
            
            data = json.dumps(message, ensure_ascii=False).encode('utf-8')
            data_with_newline = data + b'\n'
            
            # 调试：记录发送的数据
            if event in ["cancel_order", "search"]:
                logger.info(f"📤 发送 {event} 命令数据到 {client_id}: 数据长度={len(data_with_newline)}, 数据内容={data_with_newline[:200]}")
            
            client.writer.write(data_with_newline)
            await client.writer.drain()
            
            # 调试：确认发送成功
            if event in ["cancel_order", "search"]:
                logger.success(f"✅ {event} 命令已成功发送到 {client_id}，数据已刷新到TCP流")
            
            return True
        except Exception as e:
            logger.error(f"❌ 发送消息到客户端 {client_id} 失败: {e}", exc_info=True)
            return False
    
    async def broadcast_to_websocket(self, message: dict, user_id: Optional[int] = None):
        """
        广播消息到WebSocket连接（支持按user_id过滤）
        
        Args:
            message: 要广播的消息
            user_id: 如果提供，只推送给该用户的WebSocket连接
        """
        await websocket_manager.broadcast(message, user_id=user_id)
    
    async def update_heartbeat(self, client_id: str, db_session=None):
        """更新心跳时间"""
        if client_id in self.clients:
            self.clients[client_id].last_heartbeat = datetime.now()
            
            # 更新数据库
            if db_session:
                from module_admin.service.agent_client_service import AgentClientService
                async with db_session() as db:
                    try:
                        await AgentClientService.update_client_status(
                            db, client_id, {"last_heartbeat": datetime.now()}
                        )
                        await db.commit()
                    except Exception as e:
                        logger.error(f"更新心跳失败: {e}")
    
    async def _update_db_client_status(
        self,
        db_session,
        client_id: str,
        ip_address: str,
        port: int,
        status_data: dict
    ):
        """更新数据库中的客户端状态"""
        from module_admin.service.agent_client_service import AgentClientService
        
        async with db_session() as db:
            try:
                # 先尝试注册或更新客户端基本信息
                await AgentClientService.register_or_update_client(
                    db, client_id, ip_address, port
                )
                
                # 更新状态信息
                if status_data:
                    await AgentClientService.update_client_status(
                        db, client_id, status_data
                    )
                
                await db.commit()
            except Exception as e:
                logger.error(f"更新数据库客户端状态失败: {e}")
                await db.rollback()
    
    async def _mark_client_offline(self, db_session, client_id: str):
        """标记客户端离线"""
        from module_admin.service.agent_client_service import AgentClientService
        
        async with db_session() as db:
            try:
                await AgentClientService.mark_offline(db, client_id)
                await db.commit()
            except Exception as e:
                logger.error(f"标记客户端离线失败: {e}")
                await db.rollback()
    
    async def _notify_client_update(self, client_id: str):
        """通知WebSocket客户端状态更新（从数据库获取完整状态，包含中文翻译）"""
        try:
            from config.get_db import get_db
            from module_admin.service.agent_client_service import AgentClientService
            from utils.status_translation import translate_client_info
            
            # 从数据库获取完整状态信息
            async for db in get_db():
                client_info = await AgentClientService.get_client_by_client_id(db, client_id)
                if client_info:
                    client_dict = translate_client_info({
                        "id": client_info.id,
                        "client_id": client_info.client_id,
                        "client_name": client_info.client_name,
                        "ip_address": client_info.ip_address,
                        "port": client_info.port,
                        "hostname": client_info.hostname or "",
                        "client_type": client_info.client_type,
                        "online_status": client_info.online_status,
                        "work_status": client_info.work_status,
                        "browser_status": client_info.browser_status,
                        "login_status": client_info.login_status,
                        "account": client_info.account or "",
                        "account_nickname": client_info.account_nickname or "",  # 🔥 添加昵称字段
                        "default_address_id": getattr(client_info, 'default_address_id', None) or "",
                        "remark": client_info.remark or "",
                        "connected_at": client_info.connected_at.isoformat() if client_info.connected_at else "",
                        "last_heartbeat": client_info.last_heartbeat.isoformat() if client_info.last_heartbeat else "",
                    })
                    client_data = {
                        "type": "client_update",
                        "client": client_dict
                    }
                    await websocket_manager.broadcast(client_data)
                break
        except Exception as e:
            logger.error(f"通知客户端更新失败: {e}")
    
    def get_client_by_id(self, client_id: str) -> Optional[TCPClient]:
        """根据ID获取客户端"""
        return self.clients.get(client_id)
    
    def get_all_clients(self) -> List[TCPClient]:
        """获取所有客户端"""
        return list(self.clients.values())
    
    def get_client_messages(self, client_id: str, limit: int = 50) -> List[dict]:
        """获取客户端消息历史"""
        messages = self.client_messages.get(client_id, [])
        return messages[-limit:] if limit else messages
    
    def is_status_changed(self, client_id: str, new_status: dict) -> bool:
        """
        检查客户端状态是否发生变化
        
        Args:
            client_id: 客户端ID
            new_status: 新的状态数据
        
        Returns:
            True: 状态有变化，需要推送到前端
            False: 状态无变化，不需要推送
        """
        # 如果是首次收到该客户端的状态，肯定是变化
        if client_id not in self.last_status_snapshot:
            return True
        
        last_status = self.last_status_snapshot[client_id]
        
        # ✅ 使用黑名单方式：排除不需要监控的字段（主要是时间戳）
        # 其他所有字段变化都会触发推送，扩展性更好
        ignored_fields = [
            'last_heartbeat',              # 心跳时间戳（每次都变）
            # 'connected_at',                # 连接时间（不变）
            # 'last_ban_check',              # 最后封禁检查时间
            # 'last_session_refresh_time',   # 最后会话刷新时间
            # 'last_risk_time',              # 最后风险检测时间
        ]
        
        # 对比所有字段（除了黑名单中的字段）
        all_fields = set(new_status.keys()) | set(last_status.keys())
        
        for field in all_fields:
            # 跳过黑名单字段
            if field in ignored_fields:
                continue
            
            old_value = last_status.get(field)
            new_value = new_status.get(field)
            
            if old_value != new_value:
                logger.info(f"🔄 [{client_id}] 状态变化: {field} = {old_value} -> {new_value}")
                return True
        
        return False
    
    def update_status_snapshot(self, client_id: str, status: dict):
        """
        更新客户端状态快照
        
        Args:
            client_id: 客户端ID
            status: 状态数据
        """
        self.last_status_snapshot[client_id] = status.copy()
    
    async def broadcast_status_update(self, client_id: str, status_data: dict):
        """
        广播状态更新到WebSocket（仅在状态变化时）
        
        Args:
            client_id: 客户端ID
            status_data: 状态数据
        """
        # 检查状态是否变化
        if not self.is_status_changed(client_id, status_data):
            # 状态无变化，不推送
            return
        
        # 状态有变化，更新快照并广播
        self.update_status_snapshot(client_id, status_data)
        
        # 获取该 agent 分配的 user_id，只推送给对应用户
        user_id = self.get_user_id_by_client_id(client_id)
        
        await websocket_manager.broadcast({
            "type": "status_update",
            "client_id": client_id,
            "data": status_data
        }, user_id=user_id)
        
        if user_id:
            logger.debug(f" [{client_id}] 状态变化已推送给 user_id={user_id}")
        else:
            logger.debug(f" [{client_id}] 状态变化已推送（未分配用户，不推送）")
    
    async def _refresh_client_user_map(self, db_session, client_id: str):
        """
        刷新client_id -> user_id映射缓存
        
        优先级：
        1. assigned_user_ids（如果已分配，取第一个用户ID）
        2. create_by（创建者ID）
        3. None（未分配）
        """
        if not db_session:
            return
        
        try:
            from module_admin.dao.agent_client_dao import AgentClientDao
            
            async with db_session() as db:
                client_info = await AgentClientDao.get_client_by_client_id(db, client_id)
                if not client_info:
                    logger.warning(f"无法找到客户端信息: {client_id}")
                    return
                
                # ✅ 优先级1: assigned_user_id（1对1分配）
                user_id = None
                if client_info.assigned_user_id:
                    user_id = client_info.assigned_user_id
                
                # ✅ 优先级2: create_by（创建者ID，兼容旧数据）
                if user_id is None and client_info.create_by:
                    user_id = client_info.create_by
                
                # ✅ 只在映射发生变化时才记录日志，减少日志输出
                old_user_id = self.client_user_map.get(client_id)
                if old_user_id != user_id:
                    if user_id:
                        logger.info(f"✅ [{client_id}] user_id映射已更新: {old_user_id} -> {user_id}")
                    else:
                        logger.info(f"⚠️ [{client_id}] user_id映射已清除（未分配）")
                
                # 更新缓存
                self.client_user_map[client_id] = user_id
        except Exception as e:
            logger.error(f"刷新client_user_map失败: {e}", exc_info=True)
    
    def get_user_id_by_client_id(self, client_id: str) -> Optional[int]:
        """
        根据client_id获取user_id（从缓存）
        
        Returns:
            user_id: 用户ID，如果未分配则返回None
        """
        return self.client_user_map.get(client_id)
    
    async def refresh_client_user_map_by_client_ids(
        self, 
        db_session, 
        client_ids: List[str]
    ):
        """
        批量刷新client_id -> user_id映射缓存（用于Agent分配后更新）
        
        Args:
            db_session: 数据库会话工厂
            client_ids: 要刷新的client_id列表
        """
        if not db_session or not client_ids:
            return
        
        try:
            from module_admin.dao.agent_client_dao import AgentClientDao
            
            async with db_session() as db:
                for client_id in client_ids:
                    client_info = await AgentClientDao.get_client_by_client_id(db, client_id)
                    if not client_info:
                        continue
                    
                    # ✅ 优先级1: assigned_user_id（1对1分配）
                    user_id = None
                    if client_info.assigned_user_id:
                        user_id = client_info.assigned_user_id
                    
                    # ✅ 优先级2: create_by（创建者ID，兼容旧数据）
                    if user_id is None and client_info.create_by:
                        user_id = client_info.create_by
                    
                    # ✅ 只在映射发生变化时才记录日志
                    old_user_id = self.client_user_map.get(client_id)
                    if old_user_id != user_id:
                        if user_id:
                            logger.info(f"✅ [{client_id}] user_id映射已更新: {old_user_id} -> {user_id}")
                        else:
                            logger.info(f"⚠️ [{client_id}] user_id映射已清除（未分配）")
                    
                    # 更新缓存
                    self.client_user_map[client_id] = user_id
        except Exception as e:
            logger.error(f"批量刷新client_user_map失败: {e}", exc_info=True)


# 全局客户端管理器
ClientManager = EnhancedClientManager()